{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook written by [Zhedong Zheng](https://github.com/zhedongzheng)\n",
    "\n",
    "![title](img/word2vec.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/Cellar/python/3.6.5/Frameworks/Python.framework/Versions/3.6/lib/python3.6/importlib/_bootstrap.py:219: RuntimeWarning: compiletime version 3.5 of module 'tensorflow.python.framework.fast_tensor_util' does not match runtime version 3.6\n",
      "  return f(*args, **kwds)\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARAMS = {\n",
    "    'min_freq': 5,\n",
    "    'window_size': 3,\n",
    "    'n_sampled': 100,\n",
    "    'embed_dim': 200,\n",
    "    'sample_words': ['six', 'gold', 'japan', 'college'],\n",
    "    'batch_size': 1000,\n",
    "    'n_epochs': 10,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_text(text):\n",
    "    text = text.replace('\\n', ' ')\n",
    "    text = re.sub('\\s+', ' ', text).strip().lower()\n",
    "\n",
    "    words = text.split()\n",
    "    word2freq = Counter(words)\n",
    "    words = [word for word in words if word2freq[word] > PARAMS['min_freq']]\n",
    "    print(\"Total words:\", len(words))\n",
    "\n",
    "    _words = set(words)\n",
    "    PARAMS['word2idx'] = {c: i for i, c in enumerate(_words)}\n",
    "    PARAMS['idx2word'] = {i: c for i, c in enumerate(_words)}\n",
    "    PARAMS['vocab_size'] = len(PARAMS['idx2word'])\n",
    "    print('Vocabulary size:', PARAMS['vocab_size'])\n",
    "\n",
    "    indexed = [PARAMS['word2idx'][w] for w in words]\n",
    "    indexed = filter_high_freq(indexed)\n",
    "    print(\"Word preprocessing completed ...\")\n",
    "    \n",
    "    return indexed\n",
    "\n",
    "def filter_high_freq(int_words, t=1e-5, threshold=0.8):\n",
    "    int_word_counts = Counter(int_words)\n",
    "    total_count = len(int_words)\n",
    "\n",
    "    word_freqs = {w: c / total_count for w, c in int_word_counts.items()}\n",
    "    prob_drop = {w: 1 - np.sqrt(t / word_freqs[w]) for w in int_word_counts}\n",
    "    train_words = [w for w in int_words if prob_drop[w] < threshold]\n",
    "\n",
    "    return train_words\n",
    "\n",
    "def make_data(int_words):\n",
    "    x, y = [], []\n",
    "    for i in range(PARAMS['window_size'], len(int_words)-PARAMS['window_size']):\n",
    "        inputs = get_x(int_words, i)\n",
    "        x.append(inputs)\n",
    "        y.append(int_words[i])\n",
    "    return np.array(x), np.array(y)\n",
    "\n",
    "def get_x(words, idx):\n",
    "    left = idx - PARAMS['window_size']\n",
    "    right = idx + PARAMS['window_size']\n",
    "    return words[left: idx] + words[idx+1: right+1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode, params):\n",
    "    W = tf.get_variable('softmax_W', [PARAMS['vocab_size'], PARAMS['embed_dim']])\n",
    "    b = tf.get_variable('softmax_b', [PARAMS['vocab_size']])\n",
    "    E = tf.get_variable('embedding', [PARAMS['vocab_size'], PARAMS['embed_dim']])\n",
    "    embedded = tf.nn.embedding_lookup(E, features['x']) # forward activation\n",
    "    embedded = tf.reduce_mean(embedded, [1])\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        loss_op = tf.reduce_mean(tf.nn.sampled_softmax_loss(\n",
    "            weights = W,\n",
    "            biases = b,\n",
    "            labels = labels,\n",
    "            inputs = embedded,\n",
    "            num_sampled = PARAMS['n_sampled'],\n",
    "            num_classes = PARAMS['vocab_size']))\n",
    "\n",
    "        train_op = tf.train.AdamOptimizer().minimize(\n",
    "            loss_op, global_step=tf.train.get_global_step())\n",
    "        \n",
    "        return tf.estimator.EstimatorSpec(mode=mode, loss=loss_op, train_op=train_op)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        normalized_E = tf.nn.l2_normalize(E, -1)\n",
    "        sample_E = tf.nn.embedding_lookup(normalized_E, features['x'])\n",
    "        similarity = tf.matmul(sample_E, normalized_E, transpose_b=True)\n",
    "        \n",
    "        return tf.estimator.EstimatorSpec(mode, predictions=similarity)\n",
    "    \n",
    "\n",
    "def print_neighbours(similarity, top_k=5):\n",
    "    for i in range(len(PARAMS['sample_words'])):\n",
    "        neighbours = (-similarity[i]).argsort()[1:top_k+1]\n",
    "        log = 'Nearest to [%s]:' % PARAMS['sample_words'][i]\n",
    "        for k in range(top_k):\n",
    "            neighbour = PARAMS['idx2word'][neighbours[k]]\n",
    "            log = '%s %s,' % (log, neighbour)\n",
    "        print(log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total words: 885720\n",
      "Vocabulary size: 9582\n",
      "Word preprocessing completed ...\n",
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpjvjhv6s5\n",
      "INFO:tensorflow:Using config: {'_model_dir': '/var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpjvjhv6s5', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x11d6f5160>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "WARNING:tensorflow:Estimator's model_fn (<function model_fn at 0x11d6fbea0>) includes params argument, but params are not passed to Estimator.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 0 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpjvjhv6s5/model.ckpt.\n",
      "INFO:tensorflow:loss = 3.5843213, step = 1\n",
      "INFO:tensorflow:global_step/sec: 46.9561\n",
      "INFO:tensorflow:loss = 3.7559571, step = 101 (2.130 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.8487\n",
      "INFO:tensorflow:loss = 3.5779653, step = 201 (2.006 sec)\n",
      "INFO:tensorflow:global_step/sec: 50.063\n",
      "INFO:tensorflow:loss = 3.5439527, step = 301 (1.997 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.523\n",
      "INFO:tensorflow:loss = 2.9269853, step = 401 (2.019 sec)\n",
      "INFO:tensorflow:global_step/sec: 50.4256\n",
      "INFO:tensorflow:loss = 2.8054752, step = 501 (1.983 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.6669\n",
      "INFO:tensorflow:loss = 3.1015215, step = 601 (2.013 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.159\n",
      "INFO:tensorflow:loss = 3.2989767, step = 701 (2.034 sec)\n",
      "INFO:tensorflow:global_step/sec: 44.0945\n",
      "INFO:tensorflow:loss = 3.3578064, step = 801 (2.268 sec)\n",
      "INFO:tensorflow:global_step/sec: 45.8696\n",
      "INFO:tensorflow:loss = 3.2208562, step = 901 (2.180 sec)\n",
      "INFO:tensorflow:global_step/sec: 45.7902\n",
      "INFO:tensorflow:loss = 3.5523255, step = 1001 (2.184 sec)\n",
      "INFO:tensorflow:global_step/sec: 46.876\n",
      "INFO:tensorflow:loss = 3.1657982, step = 1101 (2.133 sec)\n",
      "INFO:tensorflow:global_step/sec: 47.6503\n",
      "INFO:tensorflow:loss = 3.388936, step = 1201 (2.099 sec)\n",
      "INFO:tensorflow:global_step/sec: 50.0243\n",
      "INFO:tensorflow:loss = 3.6487749, step = 1301 (1.999 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.3021\n",
      "INFO:tensorflow:loss = 2.6323447, step = 1401 (2.028 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.8968\n",
      "INFO:tensorflow:loss = 3.0210733, step = 1501 (2.004 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.5221\n",
      "INFO:tensorflow:loss = 3.1229134, step = 1601 (2.019 sec)\n",
      "INFO:tensorflow:global_step/sec: 50.957\n",
      "INFO:tensorflow:loss = 3.440561, step = 1701 (1.962 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.5455\n",
      "INFO:tensorflow:loss = 3.1347473, step = 1801 (2.018 sec)\n",
      "INFO:tensorflow:global_step/sec: 52.1922\n",
      "INFO:tensorflow:loss = 3.2595341, step = 1901 (1.916 sec)\n",
      "INFO:tensorflow:global_step/sec: 50.2213\n",
      "INFO:tensorflow:loss = 2.5702322, step = 2001 (1.991 sec)\n",
      "INFO:tensorflow:global_step/sec: 50.2411\n",
      "INFO:tensorflow:loss = 3.0401607, step = 2101 (1.990 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.3129\n",
      "INFO:tensorflow:loss = 3.030052, step = 2201 (2.028 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.6317\n",
      "INFO:tensorflow:loss = 2.6948278, step = 2301 (2.015 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.0273\n",
      "INFO:tensorflow:loss = 2.6145165, step = 2401 (2.040 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.3446\n",
      "INFO:tensorflow:loss = 2.3987436, step = 2501 (2.027 sec)\n",
      "INFO:tensorflow:global_step/sec: 49.6388\n",
      "INFO:tensorflow:loss = 2.5045497, step = 2601 (2.015 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 2698 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpjvjhv6s5/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 2.3603628.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpjvjhv6s5/model.ckpt-2698\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "Nearest to [six]: ncnb, charlotte, pioneer, modifications, involves,\n",
      "Nearest to [gold]: ounce, silver, platinum, bullion, ounces,\n",
      "Nearest to [japan]: permits, ab, against, remainder, jeopardize,\n",
      "Nearest to [college]: school, father, teaches, wonderful, done,\n"
     ]
    }
   ],
   "source": [
    "with open('../temp/ptb_train.txt') as f:\n",
    "    x_train, y_train = make_data(preprocess_text(f.read()))\n",
    "\n",
    "estimator = tf.estimator.Estimator(model_fn)\n",
    "\n",
    "estimator.train(tf.estimator.inputs.numpy_input_fn(\n",
    "    x = {'x': x_train},\n",
    "    y = np.expand_dims(y_train, -1),\n",
    "    batch_size = PARAMS['batch_size'],\n",
    "    num_epochs = PARAMS['n_epochs'],\n",
    "    shuffle = True))\n",
    "\n",
    "sim = np.array(list(estimator.predict(tf.estimator.inputs.numpy_input_fn(\n",
    "    x = {'x': np.array([PARAMS['word2idx'][w] for w in PARAMS['sample_words']])},\n",
    "    shuffle = False))))\n",
    "\n",
    "print_neighbours(sim)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
